package org.springblade.modules.mes.base.handle;

import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class MyWebSocketHandler extends TextWebSocketHandler {
	private static final Map<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();
	@Override
	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
		String userName = session.getAttributes().get("userName").toString();
		SESSIONS.put(userName, session);
		System.out.println(String.format("成功建立连接~ userName: %s", userName));
	}
	@Override
	public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
		System.out.println(message.toString());
	}
	@Override
	public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
		System.out.println("连接出错");
		if (session.isOpen()) {
			session.close();
		}
	}
	@Override
	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
		System.out.println("连接已关闭,status:" + closeStatus);
	}
	@Override
	public boolean supportsPartialMessages() {
		return false;
	}
	/**
	 * 指定发消息
	 *
	 * @param message
	 */
	public static void sendMessage(String userName, String message) {
		WebSocketSession webSocketSession = SESSIONS.get(userName);
		if (webSocketSession == null || !webSocketSession.isOpen()) return;
		try {
			webSocketSession.sendMessage(new TextMessage(message));
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	/**
	 * 群发消息
	 *
	 * @param message
	 */
	public static void fanoutMessage(String message) {
		SESSIONS.keySet().forEach(us -> sendMessage(us, message));
	}
}
